package com.wsoft.websocket.config;

import com.wsoft.core.utils.LoginUserUtil;
import com.wsoft.vo.LoginUserVO;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.http.server.ServletServerHttpRequest;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.server.support.HttpSessionHandshakeInterceptor;

import java.util.Map;

/**
 * @Author Zhang gukai
 * @Date 2025/1/9 14:40
 */
@Component
public class WebSocketHandlerInterceptor extends HttpSessionHandshakeInterceptor {

    @Override
    public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler,
                                   Map<String, Object> attributes) throws Exception {
        if (request instanceof ServletServerHttpRequest) {

            // 从请求头中获取Token
            ServletServerHttpRequest servletRequest = (ServletServerHttpRequest) request;
            String client = servletRequest.getServletRequest().getParameter("client");
            LoginUserVO loginUserVO = LoginUserUtil.getUserInfo().getLoginUserVO();
            // 放入attributes中，可以在处理器的WebSocketSession中取出
            if (loginUserVO != null) {
                attributes.put("WEBSOCKET_USERID", loginUserVO.getId());
                attributes.put("WEBSOCKET_CLIENT", client);
            }
            }
        return super.beforeHandshake(request, response, wsHandler, attributes);
    }

    @Override
    public void afterHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler,
                               Exception ex) {
        super.afterHandshake(request, response, wsHandler, ex);
    }
}
